import http.client
import json
import re
import time
import uiautomator2 as u2
from DiyAutoMessageTool import click_element
from DownloadImage import download_image  
from ChatConfig import config
class Painter:
    token=config.CZHANAI_TOKEN
    headers = {
   'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
   'Content-Type': 'application/json',
   'Auth-Token':token
}   
    model_dict={'二次元':1,'写实':15,'Q版':853,'黏土':1795,'新春佳人':974,'新春儿插':977,'儿童肖像':1389,'动漫':1811}
    def __init__(self,u2_connect,description,model_style):
        self.u2_connect=u2_connect
        self.description=description
        self.model_style=model_style
        self.painting_by_description
        self.conn = http.client.HTTPSConnection("ai.huashi6.com")
        #发送图片
    def __send_picture(self, picture_path,picture_name)->bool:
        self.u2_connect.push(picture_path, "/sdcard/"+picture_name)
        self.u2_connect.shell(f"am broadcast -a android.intent.action.MEDIA_SCANNER_SCAN_FILE -d file:///sdcard/{picture_name}")
        if click_element(self.u2_connect,resourceId="com.tencent.mm:id/bjz",description="更多功能按钮，已折叠"):
            if click_element(self.u2_connect,text="你可能要发送的照片：",resourceId="com.tencent.mm:id/llv",timeout=10):
                if click_element(self.u2_connect,text="原图",resourceId="com.tencent.mm:id/km9",timeout=10):
                    if click_element(self.u2_connect,text="发送",resourceId="com.tencent.mm:id/kaq",timeout=10):
                        return True
        return False

        #创建绘画任务
    def __create_task(self)->str:
        payload = json.dumps({
            "prompt": self.description,
            "modelStyleId": self.model_dict[self.model_style]
            })
        self.conn.request("POST", "/aiapi/v1/draw", payload, self.headers)
        res = self.conn.getresponse()
        data = res.read().decode("utf-8")
        role=re.compile('"paintingSign":"([^\n]{32})"')
        match=role.search(data)
        if(match):
            self.taskId=match.group(1)
            return data
        else:
            return None
    #查询任务进度
    def __select_task(self):
        payload = json.dumps({
        "taskId": self.taskId
        })
        self.conn.request("POST", "/aiapi/v1/task/detail", payload, self.headers)
        res = self.conn.getresponse()
        data = res.read().decode("utf-8")
        role=re.compile('"state":"([^\n]{4,8})"')
        match=role.search(data)
        if(match):
            self.state=match.group(1)
        if(self.state and self.state=="success"):
            role=re.compile('"imgUrl":"([^\n]*)"},')
            match=role.search(data)
            if(match):
                self.imgUrl=match.group(1)
            else:
                print(f"出错了,data={data}")
            return True
        return False
    # 发送图片
    def weixin_send_image(self):
        if self.__select_task():
            download_image(self.imgUrl,"./image/img.jpg")
            if self.__send_picture("./image/img.jpg",str(int(time.time()))+".jpg"):
                print(f"Info:{self.description}的图片发送成功！")
                return True
            else:
                print(f"Info:{self.description}的图片发送失败！")
                return False
        else:
            print(f"Info:正在等待{self.description}图片生成")
            return False
    # 创建AI画图任务
    def painting_by_description(self)->bool:
        try:
            if self.__create_task():
                # 此时任务已创建
                return True
            else:
                return False
        except Exception as e:
            print("Error:文生图异常报错",e)
            return False
# 测试文生图
if __name__ == "__main__":
    pass
    # alpainter=Painter(u2.connect())
    # alpainter.painting_by_description("猫咪")
    

    